import contextlib
from nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2
from nnunet.network_architecture.initialization import InitWeights_He
from nnunet.network_architecture.nyunet import (
    NYUNet,
    NYUNet01,
    NYUNet0101,
    NYUNet0102,
    NYUNet0103,
    NYUNet02,
    NYUNet03,
    NYUNet0301,
    NYUNet0302,
    NYUNet0303,
    NYUNet030301,
    NYUNet0304,
    NYUNet04,
    NYUNet05,
    NYUNet06,
    NYUNet07,
    NYUNet08,
    NYUNet09,
    NYUNet1001,
    NYUNet1002,
    NYUNet1003,
    NYUNet1004,
    NYUNet1005,
    NYUNet1006,
    NYUNet1007,
    NYUNet11,
)
import torch
from torch import nn

class nnUNetTrainerV2_yusongli(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 2


# class nnUNetTrainerV2_yangce(nnUNetTrainerV2):
#     def __init__(self, *args: object, **kwargs: object) -> None:
#         super().__init__(*args, **kwargs)
#         self.max_num_epochs = 300

#     def load_plans_file(self) -> None:
#         super().load_plans_file()
#         self.plans['plans_per_stage'][0]['batch_size'] = 30


class nYUNetTrainerV2(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000
        self.volume_shape = (24, 64, 80)
        self.model_class = None

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 30
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size  # default: 30
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size

    def initialize_network(self):
        """inference_apply_nonlin to sigmoid + larger unet"""
        if self.threeD:
            conv_op = nn.Conv3d
            dropout_op = nn.Dropout3d
            norm_op = nn.BatchNorm3d

        else:
            conv_op = nn.Conv2d
            dropout_op = nn.Dropout2d
            norm_op = nn.BatchNorm2d

        norm_op_kwargs = {'eps': 1e-5, 'affine': True}
        dropout_op_kwargs = {'p': 0, 'inplace': True}
        net_nonlin = nn.LeakyReLU
        net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        self.network = self.model_class(
            self.num_input_channels,
            self.base_num_features,
            self.num_classes,
            len(self.net_num_pool_op_kernel_sizes),
            self.conv_per_stage,
            2,
            conv_op,
            norm_op,
            norm_op_kwargs,
            dropout_op,
            dropout_op_kwargs,
            net_nonlin,
            net_nonlin_kwargs,
            True,
            False,
            lambda x: x,
            InitWeights_He(1e-2),
            self.net_num_pool_op_kernel_sizes,
            self.net_conv_kernel_sizes,
            False,
            True,
            True,
            320,
            encoder_scale=1,
            axial_attention=True,
            heads=1,
            dim_heads=4,
            volume_shape=self.volume_shape,
            no_attention=[0],
            axial_bn=True,
            sum_axial_out=True,
            residual_attention=True,
        )
        print(self.network)
        if torch.cuda.is_available():
            self.network.cuda()
        self.network.inference_apply_nonlin = nn.Sigmoid()


# ===============
# === cluster ===
# ===============
class cluster_nnunet_00(nnUNetTrainerV2_yusongli):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6


class cluster_nnunet_01(nnUNetTrainerV2_yusongli):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6


class cluster_nnunet_02(nnUNetTrainerV2_yusongli):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6


class cluster_nnunet_03(nnUNetTrainerV2_yusongli):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6


class cluster_nnunet_04(nnUNetTrainerV2_yusongli):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6


# ==================
# === caibaoshuo ===
# ==================
class new_coarse(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 2
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_coarse_foldall(new_coarse):
    pass


class new_coarse_fold0(new_coarse):
    pass


class new_coarse_fold1(new_coarse):
    pass


class new_coarse_fold2(new_coarse):
    pass


class new_coarse_fold3(new_coarse):
    pass


class new_coarse_fold4(new_coarse):
    pass


class new_coarse_00(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 30


class new_coarse_01(new_coarse_00):
    pass


class new_coarse_02(new_coarse_00):
    pass


class new_fine_nnunet(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nnunet_epoch500(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nnunet_epoch500_01(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_02(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_03(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_04(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_05(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_06(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_epoch500_07(new_fine_nnunet_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500(new_fine_nnunet_epoch500):
    pass

class new_fine_nnunet_fp32_epoch500_01(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_02(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_03(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_04(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_05(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_06(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_07(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_08(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_09(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_10(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_11(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_12(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_13(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_14(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_15(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_fp32_epoch500_16(new_fine_nnunet_fp32_epoch500):
    pass


class new_fine_nnunet_epoch1000(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nnunet_epoch1000_01(new_fine_nnunet_epoch1000):
    pass


class new_fine_nnunet_epoch1000_02(new_fine_nnunet_epoch1000):
    pass


class new_fine_nnunet_epoch1000_03(new_fine_nnunet_epoch1000):
    pass


class new_fine_nnunet_epoch2000(nnUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nnunet_epoch2000_01(new_fine_nnunet_epoch2000):
    pass


# ==============
# === NYUNet ===
# ==============

class new_fine_nyunet_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet_epoch500_01(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_02(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_03(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_04(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_05(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_06(new_fine_nyunet_epoch500):
    pass


class new_fine_nyunet_epoch500_07(new_fine_nyunet_epoch500):
    pass


# =================
# ===  NYUNet01 ===
# =================

class new_fine_nyunet01_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet01

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet01_epoch500_01(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_02(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_03(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_04(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_05(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_06(new_fine_nyunet01_epoch500):
    pass


class new_fine_nyunet01_epoch500_debug(new_fine_nyunet01_epoch500):
    pass


# ===================
# ===  NYUNet0101 ===
# ===================

class new_fine_nyunet0101_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0101

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0101_epoch500_debug(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_01(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_02(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_03(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_04(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_05(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_06(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_07(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch500_08(new_fine_nyunet0101_epoch500):
    pass


class new_fine_nyunet0101_epoch300(new_fine_nyunet0101_epoch500):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300


class new_fine_nyunet0101_epoch1000(new_fine_nyunet0101_epoch500):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000


class new_fine_nyunet0101_epoch1500(new_fine_nyunet0101_epoch500):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1500


class new_fine_nyunet0101_epoch2000(new_fine_nyunet0101_epoch500):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000


class new_fine_nyunet0101_epoch500_batch01(new_fine_nyunet0101_epoch500):
    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 1
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0101_epoch500_batch02(new_fine_nyunet0101_epoch500):
    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 2
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0101_epoch500_batch06(new_fine_nyunet0101_epoch500):
    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 6
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0102 ===
# ===================

class new_fine_nyunet0102_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0102

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0103 ===
# ===================

class new_fine_nyunet0103_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0103

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet02 ===
# =================

class new_fine_nyunet02_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet02

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet03 ===
# =================

class new_fine_nyunet03_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet03

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0301 ===
# ===================

class new_fine_nyunet0301_debug(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1
        self.model_class = NYUNet0301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0301_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0301_epoch500_01(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0302 ===
# ===================

class new_fine_nyunet0302_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0302

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet0302_epoch500_01(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0302

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0303 ===
# ===================

class new_fine_nyunet0303_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0303

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =====================
# ===  NYUNet030301 ===
# =====================

class new_fine_nyunet030301_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet030301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet030301_epoch500_01(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet030301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet030301_epoch500_02(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet030301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet030301_epoch500_03(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet030301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet030301_epoch500_04(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet030301

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# ===================
# ===  NYUNet0304 ===
# ===================

class new_fine_nyunet0304_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet0304

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet04 ===
# =================

class new_fine_nyunet04_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet04

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet05 ===
# =================

class new_fine_nyunet05_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet05

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet06 ===
# =================

class new_fine_nyunet06_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet06

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet06_epoch500_debug(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet06

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet07 ===
# =================

class new_fine_nyunet07_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet07

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet08 ===
# =================

class new_fine_nyunet08_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet08

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet09 ===
# =================

class new_fine_nyunet09_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet09

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet10 ===
# =================

class new_fine_nyunet1001_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1001

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1002_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1002

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1003_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1003

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1004_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1004

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1005_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1005

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1006_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1006

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_nyunet1007_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet1007

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


# =================
# ===  NYUNet11 ===
# =================

class new_fine_nyunet11_epoch500(nYUNetTrainerV2):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500
        self.model_class = NYUNet11

    def load_plans_file(self) -> None:
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size
